#!/bin/env python
import numpy as np
import os
import re
import argparse

HRindexfile = "HuangRhys.dat"
Fccfile = "fcclasses.out"

parser = argparse.ArgumentParser(description="对于 fcclass3 计算得到的 HR index 数据进行处理，需要当前目录下的 HuangRhys.dat fcclasses.out")
args = parser.parse_args()


def ReadHRindex(HRindexfile):
    print(f"reading ReadHRindex:{HRindexfile}")
    HRindex = []
    slag = 0
    RH = open(HRindexfile, 'r', encoding='utf8')
    for line in RH:
        if 'Huang-Rhys' in line:
            slag = 1
            continue
        if slag == 1 and len(line.strip().split()) ==2:
            HRindex.append(line.strip().split()[-1])
    RH.close()
    return HRindex
if not os.path.isfile(HRindexfile):
    print(f"File {HRindexfile} does not exist.")
    pass
else:
    HRindex = np.array(ReadHRindex(HRindexfile))

def Readfreq(Fccfile):
    print(f"reading Readfcc.out:{Fccfile}\n")
    pattern = r'------------------------\s+FREQUENCIES \(cm-1\)\s+------------------------\s+((?:\s*\d+\s+\d+\.\d+\s*)+)\s+------------------------'
    FR = open(Fccfile, 'r', encoding='utf8')
    content = FR.read()
    matches = re.findall(pattern, content)
    all_frequencies = []
    for match in matches:
        frequencies = re.findall(r'\d+\s+(\d+\.\d+)', match)
        all_frequencies.append(frequencies)
    return all_frequencies

if not os.path.isfile(Fccfile):
    print(f"File {Fccfile} does not exist.")
    pass
else:
    all_frequencies = np.array(Readfreq(Fccfile)).astype(np.float32)

with open('freq-hrindex.txt', 'w') as f:
    f.write("Vibration frequency (cm-1)\t Huang-Rhys index \n")  # 写入表头信息
    for hr, freq in zip(all_frequencies[0],HRindex):
        f.write(f"{hr}\t {freq}\n")  # 将 HRindex 和 all_frequencies 中的元素写入到文件，用制表符分隔并换行

print("数据已写入到文件 'freq-hrindex.txt' 中。")
